/**
 * Copyright (c) Huawei Technologies Co., Ltd. 2022. All Rights Reserved.
 * SPDX-License-Identifier: Apache-2.0.
 */

#include <sys/stat.h>
#include <fcntl.h>
#include <unistd.h>
#include <string.h>
#include <gnutls/x509.h>
#include <gnutls/gnutls.h>
#include <nettle/sha2.h>
#include <cjson/cJSON.h>

#include "enclave.h"
#include "util.h"

#define ENCLVA_MAGIC            0x20210401
#define ENCLVA_VERSION          1
#define UCOMP_VMLINUX           "/usr/local/share/qingtian/enclave/vmlinux.bin"
#define ALIGN_BLOCK             0x1000
#define ALIGN_MASK              (ALIGN_BLOCK - 1)
#define BLOCK_ALIGN(size)       ((size + ALIGN_MASK) & (~ALIGN_MASK))
#define SIGNATURE_MAX_SIZE      512

static uint8_t pcr0[SHA384_DIGEST_SIZE] = { 0 };
static uint8_t pcr8[SHA384_DIGEST_SIZE] = { 0 };

enum PCR_TYPE {
    PCR0 = 0,
    PCR8
};

static int64_t get_file_size(const char *filename)
{
    struct stat s;

    if (stat(filename, &s)) {
        qt_log_err(QT_FILE_OPS_ERROR, QT_DISPLAY_ADDITIONAL_ERROR,
                   "Operation: stat, File: %s, errno: %d [%m])", filename, errno);
        return -1;
    }

    if (!S_ISREG(s.st_mode) && !S_ISBLK(s.st_mode)) {
        qt_log_err(QT_INVALID_ARGS_ERROR, QT_DISPLAY_ADDITIONAL_ERROR,
                   "%s is not a file or bdev", filename);
        return -1;
    }

    return s.st_size;
}

static int img_write(int fd, void *buf, size_t count, off_t offset)
{
    int ret;

    while (count > 0) {
        ret = pwrite(fd, buf, count, offset);
        if (ret < 0) {
            if (errno == EINTR)
                continue;

            qt_log_err(QT_UNDEINED_ERROR_NUM, QT_NOT_DISPLAY_ADDITIONAL_ERROR,
                       "Write offset %ld failed(%d)", offset, errno);
            return errno;
        }
        count -= ret;
    }

    return 0;
}

static int img_read(int fd, void *buf, size_t count, off_t offset)
{
    int ret;

    while (count > 0) {
        ret = pread(fd, buf, count, offset);
        if (ret < 0) {
            if (errno == EINTR)
                continue;

            qt_log_err(QT_FILE_OPS_ERROR, QT_DISPLAY_ADDITIONAL_ERROR,
                       "Operations: pread offset %ld failed, errno: %d [%m]", offset, errno);
            return errno;
        }
        count -= ret;
    }

    return 0;
}

static void get_pcr(uint8_t *img_ptr, struct enclave_img_hdr *hdr, int type)
{
    struct sha384_ctx ctx;
    uint8_t zero[SHA384_DIGEST_SIZE] = { 0 };
    uint8_t digest[SHA384_DIGEST_SIZE];
    uint8_t *result = NULL;

    sha384_init(&ctx);

    /* steps calculate PCR
     * 1, digest = sha384(data)
     * 2, pcr = sha384( 0 || data)
     */

    switch (type) {
        case PCR0:
            sha384_update(&ctx, hdr->cmdline_size, img_ptr + hdr->cmdline_offset);
            sha384_update(&ctx, hdr->kernel_size, img_ptr + hdr->kernel_offset);
            sha384_update(&ctx, hdr->initrd_size, img_ptr + hdr->initrd_offset);
            result = pcr0;
            break;
        case PCR8:
            sha384_update(&ctx, hdr->certificate_size, img_ptr + hdr->certificate_offset);
            result = pcr8;
            break;
        default:
            qt_log_err(QT_UNDEINED_ERROR_NUM, QT_NOT_DISPLAY_ADDITIONAL_ERROR,
                       "not support type: %d", type);
            return;
    }

    sha384_digest(&ctx, SHA384_DIGEST_SIZE, digest);

    /* sha384_digest will init ctx again */
    sha384_update(&ctx, SHA384_DIGEST_SIZE, zero);
    sha384_update(&ctx, SHA384_DIGEST_SIZE, digest);
    sha384_digest(&ctx, SHA384_DIGEST_SIZE, result);
    return;
}

static int calc_pcrs(int fd, struct enclave_img_hdr *hdr)
{
    uint8_t *img_ptr = NULL;

    img_ptr = mmap(NULL, hdr->total_size, PROT_READ, MAP_PRIVATE, fd, 0);
    if (img_ptr == MAP_FAILED) {
        qt_log_err(QT_MMAP_MEM_ERROR, QT_DISPLAY_ADDITIONAL_ERROR,
                   "Mmap enclave image failed, errno: %d [%m]", errno);
        return -1;
    }

    get_pcr(img_ptr, hdr, PCR0);

    if (hdr->certificate_size != 0) {
        get_pcr(img_ptr, hdr, PCR8);
    }

    munmap(img_ptr, hdr->total_size);
    return 0;
}

static void print_pcrs(void)
{
    int i;
    char buf[SHA384_DIGEST_SIZE * 2 + 1];
    cJSON *monitor;
    char *console_str = NULL;

    monitor = cJSON_CreateObject();
    if (monitor == NULL) {
        qt_log_err(QT_CREATE_CJSON_OBJ_ERROR, QT_DISPLAY_ADDITIONAL_ERROR,
                   "Please check cjson library");
        return;
    }

    if (cJSON_AddStringToObject(monitor, "digest", "SHA384") == NULL) {
        qt_log_err(QT_ADD_STR_TO_CJSON_OBJ_ERROR, QT_DISPLAY_ADDITIONAL_ERROR,
                   "Please check cjson library");
        goto end;
    }

    log_info("digest: SHA384");
    for (i = 0; i < SHA384_DIGEST_SIZE; i++)
        sprintf(buf + 2 * i, "%02x", pcr0[i]);
    log_info("PCR0: %s", buf);
    if (cJSON_AddStringToObject(monitor, "PCR0", buf) == NULL) {
        qt_log_err(QT_ADD_STR_TO_CJSON_OBJ_ERROR, QT_DISPLAY_ADDITIONAL_ERROR,
                   "Add PCR0 in json object fail");
        goto end;
    }

    for (i = 0; i < SHA384_DIGEST_SIZE; i++)
        sprintf(buf + 2 * i, "%02x", pcr8[i]);
    if (cJSON_AddStringToObject(monitor, "PCR8", buf) == NULL) {
        qt_log_err(QT_ADD_STR_TO_CJSON_OBJ_ERROR, QT_DISPLAY_ADDITIONAL_ERROR,
                   "Add PCR8 in json object fail");
        goto end;
    }
    log_info("PCR8: %s", buf);

    console_str = cJSON_Print(monitor);
    if (console_str == NULL) {
        qt_log_err(QT_GET_JSON_OBJ_ERROR, QT_DISPLAY_ADDITIONAL_ERROR,
                   "Please check cjson library");
        goto end;
    }

    printf("%s\n", console_str);

    free(console_str);
end:
    cJSON_Delete(monitor);
    return;
}

/* hdr(4096) | cmdline(4096) | kernel | initrd | certificate | pcr0_signature(4096) */
static int calc_img_layout(struct enclave_img_info *info,
                           struct enclave_img_hdr *hdr)
{
    int64_t size;

    hdr->magic = ENCLVA_MAGIC;
    hdr->version = ENCLVA_VERSION;

    hdr->total_size = BLOCK_ALIGN(sizeof(*hdr)) + ALIGN_BLOCK;
    hdr->cmdline_offset = BLOCK_ALIGN(sizeof(*hdr));
    hdr->cmdline_size = info->cmdline ? strlen(info->cmdline) : 0;

    hdr->kernel_offset = hdr->total_size;
    size = get_file_size(info->kernel);
    if (size < 0)
        return -1;
    hdr->kernel_size = size;
    hdr->total_size += BLOCK_ALIGN(size);

    size = get_file_size(UCOMP_VMLINUX);
    if (size < 0)
        return -1;
    hdr->uncomp_kernel_size = size;

    hdr->initrd_offset = hdr->total_size;
    size = get_file_size(info->initrd);
    if (size < 0)
        return -1;
    hdr->initrd_size = size;
    hdr->total_size += BLOCK_ALIGN(size);

    if (!info->private_key || !info->signing_certificate)
    return  0;

    hdr->certificate_offset = hdr->total_size;
    size = get_file_size(info->signing_certificate);
    if (size < 0)
        return -1;
    hdr->certificate_size = size;
    hdr->total_size += BLOCK_ALIGN(size);

    hdr->signature_offset = hdr->total_size;
    hdr->signature_size = 0;
    hdr->total_size += ALIGN_BLOCK;

    return 0;
}

static int image_create(const char *file, int64_t size)
{
    int fd;

    fd = open(file, O_RDWR | O_CREAT | O_TRUNC, S_IRUSR | S_IWUSR | S_IRGRP);
    if (!fd) {
        qt_log_err(QT_FILE_OPS_ERROR, QT_DISPLAY_ADDITIONAL_ERROR,
                   "Operations: open, File: %s, errno:%d [%m]", file, errno);
        return -1;
    }

    if (ftruncate(fd, size)) {
        qt_log_err(QT_FILE_OPS_ERROR, QT_DISPLAY_ADDITIONAL_ERROR,
                   "Operations: truncate, File %s, errno: %d [%m]", file, errno);
        close(fd);
        return -1;
    }

    return fd;
}

static int img_build(int img_fd, char *file, int64_t total, int64_t start)
{
    int fd;
    int size, off = 0;
    char buf[ALIGN_BLOCK];

    if (total <= 0 || !file)
        return 0;

    fd = open(file, O_RDONLY, S_IRUSR | S_IWUSR | S_IRGRP);
    if (fd < 0) {
        qt_log_err(QT_FILE_OPS_ERROR,
                   QT_DISPLAY_ADDITIONAL_ERROR,
                   "Operations: open, File: %s, errno: %d [%m]", file, errno);
        return fd;
    }

    while (total > 0) {
        size = MIN(total, ALIGN_BLOCK);
        if (img_read(fd, buf, size, off)) {
            close(fd);
            return -1;
        }

        if (img_write(img_fd, buf, size, start + off)) {
            close(fd);
            return -1;
        }
        off += size;
        total -= size;
    }

    close(fd);
    return 0;
}

static int get_private_key_data(char *private_key, gnutls_datum_t *dt)
{
    int64_t size;
    int priv_fd;
    uint8_t *private_ptr = NULL;

    size = get_file_size(private_key);
    if (size < 0)
        return -1;

    priv_fd = open(private_key, O_RDONLY);
    if (priv_fd < 0) {
        qt_log_err(QT_FILE_OPS_ERROR, QT_DISPLAY_ADDITIONAL_ERROR,
                   "Operations: open, File: %s, errno: %d [%m]", private_key, errno);
        return -1;
    }

    private_ptr = mmap(NULL, size, PROT_READ, MAP_PRIVATE, priv_fd, 0);
    if (!private_ptr) {
        qt_log_err(QT_MMAP_MEM_ERROR, QT_DISPLAY_ADDITIONAL_ERROR,
                   "Mmap %s failed: [%m]", private_key);
        close(priv_fd);
        return -1;
    }

    dt->data = private_ptr;
    dt->size = size;

    close(priv_fd);
    return 0;
}

static int img_add_signature(int fd, char *private_key, struct enclave_img_hdr *hdr)
{
    int ret;
    int retv = -1;
    uint8_t signature_data[SIGNATURE_MAX_SIZE];
    size_t signature_size = sizeof(signature_data);
    gnutls_datum_t privkey_dt, pcr_dt;
    gnutls_x509_privkey_t privkey;

    ret = get_private_key_data(private_key, &privkey_dt);
    if (ret < 0) {
        qt_log_err(QT_UNDEINED_ERROR_NUM, QT_NOT_DISPLAY_ADDITIONAL_ERROR,
                   "Get privkey data failed");
        return retv;
    }

    ret = gnutls_x509_privkey_init(&privkey);
    if (ret < 0) {
        qt_log_err(QT_UNDEINED_ERROR_NUM, QT_NOT_DISPLAY_ADDITIONAL_ERROR,
                   "Priv key init failed");
        goto out_key;
    }

    ret = gnutls_x509_privkey_import(privkey, &privkey_dt, GNUTLS_X509_FMT_PEM);
    if (ret < 0) {
        qt_log_err(QT_UNDEINED_ERROR_NUM, QT_NOT_DISPLAY_ADDITIONAL_ERROR,
                   "Import privkey failed %d", ret);
        goto out;
    }

    pcr_dt.data = pcr0;
    pcr_dt.size = SHA384_DIGEST_SIZE;
    ret = gnutls_x509_privkey_sign_data(privkey, GNUTLS_DIG_SHA384, 0, &pcr_dt,
                                        signature_data, &signature_size);
    if (ret < 0) {
        qt_log_err(QT_UNDEINED_ERROR_NUM, QT_NOT_DISPLAY_ADDITIONAL_ERROR,
                   "Sign data failed %d", ret);
        goto out;
    }

    if (img_write(fd, signature_data, signature_size, hdr->signature_offset)) {
        qt_log_err(QT_UNDEINED_ERROR_NUM, QT_NOT_DISPLAY_ADDITIONAL_ERROR,
                   "Write signature data failed");
        goto out;
    }

    /* update signature_size at header */
    hdr->signature_size = signature_size;
    if (img_write(fd, hdr, sizeof(*hdr), 0)) {
        qt_log_err(QT_UNDEINED_ERROR_NUM, QT_NOT_DISPLAY_ADDITIONAL_ERROR,
                   "update signature_size failed");
        goto out;
    }

    retv = 0;
out:
    gnutls_x509_privkey_deinit(privkey);
out_key:
    munmap(privkey_dt.data, privkey_dt.size);

    return retv;
}

static int enclave_img_info_check(struct enclave_img_info *info)
{
    if (!info || !info->initrd || !info->kernel || !info->output) {
        qt_log_err(QT_MISSING_ARGS_ERROR, QT_DISPLAY_ADDITIONAL_ERROR,
                   "Missing enclave_img_info parameters");
        return -1;
    }

    if (!info->initrd) {
        qt_log_err(QT_MISSING_ARGS_ERROR, QT_DISPLAY_ADDITIONAL_ERROR,
                   "Missing initrd parameters");
        return -1;
    }

    if (!info->kernel) {
        qt_log_err(QT_MISSING_ARGS_ERROR, QT_DISPLAY_ADDITIONAL_ERROR,
                   "Missing kernel parameters");
        return -1;
    }

    if (!info->output) {
        qt_log_err(QT_MISSING_ARGS_ERROR, QT_DISPLAY_ADDITIONAL_ERROR,
                   "Missing outputs parameters");
        return -1;
    }

    if (info->cmdline && strlen(info->cmdline) >= ALIGN_BLOCK) {
        qt_log_err(QT_IMG_INFO_CHECK_ERROR, QT_DISPLAY_ADDITIONAL_ERROR,
                   "CMDline is too long!");
        return -1;
    }

    return 0;
}

int enclave_img_build(struct enclave_img_info *info)
{
    int fd;
    int ret = -1;
    struct enclave_img_hdr hdr = { 0 };

    if (enclave_img_info_check(info) < 0)
        return -1;

    if (calc_img_layout(info, &hdr))
        return -1;

    fd = image_create(info->output, hdr.total_size);
    if (fd < 0)
        return -1;

    /* write hdr */
    if (img_write(fd, &hdr, sizeof(hdr), 0)) {
        qt_log_err(QT_WRITE_IMG_HDR_ERROR, QT_DISPLAY_ADDITIONAL_ERROR,
                   "Please check the log for details");
        goto out;
    }

    if (info->cmdline &&
        img_write(fd, info->cmdline, hdr.cmdline_size,
              hdr.cmdline_offset)) {
        qt_log_err(QT_WRITE_CMDLINE_ERROR, QT_DISPLAY_ADDITIONAL_ERROR,
                   "Please check the log for details");
        goto out;
    }

    if (img_build(fd, info->kernel, hdr.kernel_size, hdr.kernel_offset)) {
        qt_log_err(QT_WRITE_KERNEL_ERROR, QT_DISPLAY_ADDITIONAL_ERROR,
                   "Please check the log for details");
        goto out;
    }

    if (img_build(fd, info->initrd, hdr.initrd_size, hdr.initrd_offset)) {
        qt_log_err(QT_WRITE_INITTRD_ERROR, QT_DISPLAY_ADDITIONAL_ERROR,
                   "Please check the log for details");
        goto out;
    }

    if (info->private_key &&
        info->signing_certificate &&
        img_build(fd, info->signing_certificate, hdr.certificate_size,
                  hdr.certificate_offset)) {
        qt_log_err(QT_WRITE_CERT_ERROR, QT_DISPLAY_ADDITIONAL_ERROR,
                   "Please check the log for details");
        goto out;
    }

    if (calc_pcrs(fd, &hdr)) {
        qt_log_err(QT_GET_PCR_ERROR, QT_DISPLAY_ADDITIONAL_ERROR,
                   "Please check the log for details");
        goto out;
    }

    if (info->private_key &&
        img_add_signature(fd, info->private_key, &hdr)) {
        qt_log_err(QT_ADD_SIG_ERROR, QT_DISPLAY_ADDITIONAL_ERROR,
                   "Please check the log for details");
        goto out;
    }

    print_pcrs();
    (void)fsync(fd);
    ret = 0;
out:
    close(fd);
    if (ret)
        unlink(info->output);

    return ret;
}
